# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: jshin

import logging
from typing import List, Union
import json

from tornado import gen

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase


class BithumbSubscriber(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._symbol = to_list(symbol)

    self._register_url(
        url="wss://pubwss.bithumb.com/pub/ws",
        headers={
            'Host':
                'pubwss.bithumb.com',
            'User-Agent':
                'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/73.0.3683.86 Safari/537.36',  # noqa: E501
        })
    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_depth = topic_map.search_topic_id('bithumb_ws_depth_%s%s' % (group_id, worker_id))[0]
    topic_trade = topic_map.search_topic_id('bithumb_ws_trade_%s%s' % (group_id, worker_id))[0]
    topic_ticker = topic_map.search_topic_id('bithumb_ws_ticker_%s%s' % (group_id, worker_id))[0]
    topic_quote = topic_map.search_topic_id('bithumb_ws_quote_%s%s' % (group_id, worker_id))[0]

    self._publisher = {
        'orderbookdepth': queue.get_writer(topic_id=topic_depth[1]),
        'transaction': queue.get_writer(topic_id=topic_trade[1]),
        'ticker': queue.get_writer(topic_id=topic_ticker[1]),
        'quote': queue.get_writer(topic_id=topic_quote[1]),
        'chrt': None,
        'control': None,
    }

    self._enable_topic_event_publisher(queue,
                                       'bithumb_ws_topic_event', [topic_depth[0], topic_trade[0]])
    self._add_teardown_callback(self._teardown_publisher)

  @gen.coroutine
  def on_open(self):
    self._ioloop.add_callback(self._send_subscribe_request)
    self._publish_topic_event('START')

  @gen.coroutine
  def _send_subscribe_request(self):
    self._logger.debug('Subscribing %s' % self._symbol)
    symbols = ["%s_KRW" % symbol for symbol in self._symbol]
    msg = json.dumps({'type': 'transaction', 'symbols': symbols})
    yield self._ws.write_message(msg)
    yield gen.sleep(0.5)  # Sleep 0.5 second before next request
    msg = json.dumps({'type': 'orderbookdepth', 'symbols': symbols})
    yield self._ws.write_message(msg)
    yield gen.sleep(0.5)  # Sleep 0.5 second before next request

    self._publish_topic_event('START')
    self._logger.info('Subscribe done')

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        if publisher is not None:
          publisher.close()
      self._publisher = None

  def on_message(self, timestamp: int, message: dict):
    if self.stopped:
      return

    if "status" in message:
      status = message["status"]
      print(timestamp, status, message)
      return

    event_type = message['type']
    try:
      writer = self._publisher[event_type]
      if writer is not None:
        writer.write(timestamp, message)
    except KeyError:
      self._logger.error('Unknown event_type: %s' % event_type)
